import sys,os
import pyfaidx


def get_contig_fasta(infasta, inagp):
    agp = open(inagp, "r")
    fasta = pyfaidx.Fasta(infasta)
    marker = 1
    hap1fasta = open("hap1.fasta", "w")
    hap2fasta = open("hap2.fasta", "w")
    hap1agp = open("hap1.agp", "w")
    hap2agp = open("hap2.agp", "w")
    hap1contigindex = 0
    hap2contigindex = 0
    for i in agp:
        tmp = i.strip().split("\t")
        scaffold = tmp[0]
        # start = int(tmp[1])
        # end = int(tmp[2])
        if i.startswith("scaffold_3639"):
            marker = 2
        contigname = tmp[5]
        if contigname == "100":
            if marker == 1:
                hap1agp.write(i)
            else:
                hap2agp.write(i)
            continue
        start = int(tmp[6]) - 1
        end = int(tmp[7])
        seq = fasta[contigname][start:end]
        contiglen = end - start
        
        if marker == 1:
            contigname = f"hap1_{hap1contigindex}"
            agpoutline = "\t".join(tmp[0:5]) + "\t" + contigname + "\t" + "1" + "\t" + str(contiglen) + "\t" + tmp[8] + "\n"
            hap1fasta.write(">" + contigname + "\n" + str(seq) + "\n")
            hap1agp.write(agpoutline)
            hap1contigindex += 1
            
        else:
            contigname = f"hap2_{hap2contigindex}"
            agpoutline = "\t".join(tmp[0:5]) + "\t" + contigname + "\t" + "1" + "\t" + str(contiglen) + "\t" + tmp[8] + "\n"
            hap2fasta.write(">" + contigname + "\n" + str(seq) + "\n")
            hap2agp.write(agpoutline)
            hap2contigindex += 1

if __name__ == "__main__":
    infasta = sys.argv[1]
    inagp = sys.argv[2]
    get_contig_fasta(infasta, inagp)
            
            
            
            
    